# %% 
# import sys
# sys.path.append('.')
from pathlib import Path
import json
from torchvision.io import read_image, ImageReadMode
from torchvision.utils import make_grid
from PIL import Image
import matplotlib.pyplot as plt

# %% 
rank_key = 'lpips' # l1, l2, lpips
rank_idx=3 # rank x
overide_seq_idx=17 # 直接指定序列的 idx , 不排序

root_dir='samples/ds_sample/rk45/2reflow_v4_s156250'
root_dir=Path(root_dir)

comp_root_dirs=[
    # 'samples/ds_sample/euler25/2reflow_v4_s30000', 
    # 'samples/ds_sample/euler25/2reflow_v3_s75000', 
    # 'samples/ds_sample/euler25/2reflow_v3_s100000', 
    # 'samples/ds_sample/euler25/2reflow_v3_s125000', 
    # 'samples/ds_sample/euler25/2reflow_v3_s150000', 
    # 'samples/ds_sample/euler25/2reflow_v3_s156250',
    ]
comp_root_dirs=[Path(p) for p in comp_root_dirs]

loss_info=json.load((root_dir / 'loss_info.json').open('r'))
comp_loss_info=[json.load((p / 'loss_info.json').open('r')) for p in comp_root_dirs]

prompts = (root_dir / 'caption.txt').open('r').read().splitlines()


# swap=False
# if swap:
#     root_dir, comp_root_dir = comp_root_dir, root_dir

print(f'root dir:\n{root_dir}\ncompany root dir:\n{comp_root_dirs}')

# %% 
def argsort(seq, key_func):
    indices = sorted(range(len(seq)), key=key_func)
    return indices
sorted_indices = argsort(loss_info,lambda x:loss_info[x][rank_key])


# %% 
img_suffix='jpg'
print(f'rank with {rank_key}')
seq_idx=sorted_indices[rank_idx]
if overide_seq_idx is not None:
    seq_idx=overide_seq_idx

# # overide seq_idx
# seq_idx=17

print(f'original index : {seq_idx}')
print(prompts[seq_idx])
print(loss_info[seq_idx])
for cli in comp_loss_info:
    print(cli[seq_idx])
    

latent = read_image(str(root_dir / 'latent' / f'latent_{seq_idx}.{img_suffix}'), mode=ImageReadMode.RGB) 
sample = read_image(str(root_dir / 'samples' / f'sample_{seq_idx}.{img_suffix}'), mode=ImageReadMode.RGB) 
samples_comp = [read_image(str(p / 'samples' / f'sample_{seq_idx}.{img_suffix}'), mode=ImageReadMode.RGB)  for p in comp_root_dirs]



grid=make_grid([latent,sample,]+samples_comp, nrow=2+len(samples_comp))
pil_grid=Image.fromarray(grid.permute(1,2,0).numpy())
print(pil_grid)
# plt.imshow(pil_grid)
# plt.show()
# pil_grid.save(f'tmp/grid.{img_suffix}')
pil_grid

# %%
